#ifndef IMAGE_HPP
#define IMAGE_HPP

#include "../math/tvector.hpp"
#include "../math/color_operator.hpp"
#include "../misc/timer.hpp"
#include "image_descript.hpp"
#include <functional>
#include <string>

#include <stb_image_write.h>

using namespace ColorOP;

class common_image_writer {
protected:
  std::string fileName;
  ~common_image_writer() {}

public:
  virtual void excute() = 0;

  common_image_writer( const i8 *fileName ) : fileName( fileName ) {}
  common_image_writer() = default;
};

template <bool alpha>
struct image_alpha_selector {};

template <>
struct image_alpha_selector<true> {
  static const ui8 _comp = 4;
};

template <>
struct image_alpha_selector<false> {
  static const ui8 _comp = 3;
};

// TODO: this code is not efficient
template <ui32 w, ui32 h, bool alpha>
class image_writer : public common_image_writer {
public:
  using generateFunc = std::function<vec4f( ui32 x, ui32 y, ui32 w, ui32 h )>;

private:
  image_descriptor<w, h, image_alpha_selector<alpha>::_comp> _desc;
  generateFunc _func;

public:
  image_writer( const i8 *fileName, const generateFunc &func )
      : _func( func ), common_image_writer( fileName ), _desc() {}

  void excute() override {
    ui32 tmp = 0;
    Timer t;
    for ( i32 y = h; y > 0; y-- ) {
      for ( i32 x = 0; x < w; x++ ) {
        vec4f color = _func( x, y, w, h );
        ui32 col = encodeRGBAF( color, true );
        _desc.getRawData()[tmp++] = getSingleChannel( col, ColorChannel::R );
        _desc.getRawData()[tmp++] = getSingleChannel( col, ColorChannel::G );
        _desc.getRawData()[tmp++] = getSingleChannel( col, ColorChannel::B );
        _desc.getRawData()[tmp++] = getSingleChannel( col, ColorChannel::A );
      }
    }
    std::cout << "Time Consume: " << t.elapsed() << "s\n";
  }

  void save() {
    if ( stbi_write_png( fileName.c_str(), w, h,
                         image_alpha_selector<alpha>::_comp,
                         (const void *)_desc.getRawData(), 0 ) == 0 )
      std::cout << "Write image failed!\n";
  }

  void excuteAndSave() { // just a simple alias for execute and save
    excute();
    save();
  }
};

#endif